import torch

import torch.nn.functional as F

from util.misc import (accuracy, get_world_size,
                       is_dist_avail_and_initialized)
import numpy as np
from queue import Queue
import math
from  .setcriterionhoi import SetCriterionHOI as setcrithoi


class SetCriterionHOI(setcrithoi):
    def __init__(self, args, matcher, weight_dict, losses, loss_type='ce_loss'):
        super(SetCriterionHOI, self).__init__(args, matcher, weight_dict, losses, loss_type)
        self.alpha = args.focal_alpha

        if args.dataset_file == 'hico':
            self.obj_nums_init = [1811, 9462, 2415, 7249, 1665, 3587, 1396, 1086, 10369, 800, \
                                  287, 77, 332, 2352, 974, 470, 1386, 4889, 1675, 1131, \
                                  1642, 185, 92, 717, 2228, 4396, 275, 1236, 1447, 1207, \
                                  2949, 2622, 1689, 2345, 1863, 408, 5594, 1178, 562, 1479, \
                                  988, 1057, 419, 1451, 504, 177, 1358, 429, 448, 186, \
                                  121, 441, 735, 706, 868, 1238, 1838, 1224, 262, 517, \
                                  5787, 200, 529, 1337, 146, 272, 417, 1277, 31, 213, \
                                  7, 102, 102, 2424, 606, 215, 509, 529, 102, 572]
        elif args.dataset_file == 'vcoco':
            self.obj_nums_init = [5397, 238, 332, 321, 5, 6, 45, 90, 59, 20, \
                                  13, 5, 6, 313, 28, 25, 46, 277, 20, 16, \
                                  154, 0, 7, 13, 356, 191, 458, 66, 337, 1364, \
                                  1382, 958, 1166, 68, 258, 221, 1317, 1428, 759, 201, \
                                  190, 444, 274, 587, 124, 107, 102, 37, 226, 16, \
                                  30, 22, 187, 320, 222, 465, 893, 213, 56, 322, \
                                  306, 13, 55, 834, 23, 104, 38, 861, 11, 27, \
                                  0, 16, 22, 405, 50, 14, 145, 63, 9, 11]
        else:
            raise

        self.obj_nums_init.append(3 * sum(self.obj_nums_init))  # 3 times fg for bg init
        if args.dataset_file == 'hico':
            self.verb_nums_init = [67, 43, 157, 321, 664, 50, 232, 28, 5342, 414, \
                                   49, 105, 26, 78, 157, 408, 358, 129, 121, 131, \
                                   275, 1309, 3, 799, 2338, 128, 633, 79, 435, 1, \
                                   905, 19, 319, 47, 816, 234, 17958, 52, 97, 648, \
                                   61, 1430, 13, 1862, 299, 123, 52, 328, 121, 752, \
                                   111, 30, 293, 6, 193, 32, 4, 15421, 795, 82, \
                                   30, 10, 149, 24, 59, 504, 57, 339, 62, 38, \
                                   472, 128, 672, 1506, 16, 275, 16092, 757, 530, 380, \
                                   132, 68, 20, 111, 2, 160, 3209, 12246, 5, 44, \
                                   18, 7, 5, 4815, 1302, 69, 37, 25, 5048, 424, \
                                   1, 235, 150, 131, 383, 72, 76, 139, 258, 464, \
                                   872, 360, 1917, 1, 3775, 1206, 1]
        elif args.dataset_file == 'vcoco':
            self.verb_nums_init = [4001, 4598, 1989, 488, 656, 3825, 367, 367, 677, 677, \
                                   700, 471, 354, 498, 300, 313, 300, 300, 622, 458, \
                                   500, 498, 489, 1545, 133, 142, 38, 116, 388]
        else:
            raise

        self.verb_nums_init.append(3 * sum(self.verb_nums_init))

        self.obj_reweight = args.obj_reweight
        self.verb_reweight = args.verb_reweight
        self.use_static_weights = args.use_static_weights
        
#########################  reweight  #####################################################
        Maxsize = args.queue_size

        if self.obj_reweight:
            self.q_obj = Queue(maxsize=Maxsize)
            self.p_obj = args.p_obj
            self.obj_weights_init = self.cal_weights(self.obj_nums_init, p=self.p_obj)

        if self.verb_reweight:
            self.q_verb = Queue(maxsize=Maxsize)
            self.p_verb = args.p_verb
            self.verb_weights_init = self.cal_weights(self.verb_nums_init, p=self.p_verb)
#########################  reweight  #####################################################


    def cal_weights(self, label_nums, p=0.5):
        num_fgs = len(label_nums[:-1])
        weight = [0] * (num_fgs + 1)
        num_all = sum(label_nums[:-1])

        for index in range(num_fgs):
            if label_nums[index] == 0: continue
            weight[index] = np.power(num_all/label_nums[index], p)

        weight = np.array(weight)
        weight = weight / np.mean(weight[weight>0])

        weight[-1] = np.power(num_all/label_nums[-1], p) if label_nums[-1] != 0 else 0

        weight = torch.FloatTensor(weight).cuda()
        return weight


    def obj_reweight_func(self, target_classes):
        if not self.obj_reweight:
            obj_weights = self.empty_weight
        elif self.use_static_weights:
            obj_weights = self.obj_weights_init
        else:
            obj_label_nums_in_batch = [0] * (self.num_obj_classes + 1)
            for target_class in target_classes:
                for label in target_class:
                    obj_label_nums_in_batch[label] += 1

            if self.q_obj.full(): self.q_obj.get()
            self.q_obj.put(np.array(obj_label_nums_in_batch))
            accumulated_obj_label_nums = np.sum(self.q_obj.queue, axis=0)
            obj_weights = self.cal_weights(accumulated_obj_label_nums, p=self.p_obj)

            aphal = min(math.pow(0.999, self.q_obj.qsize()),0.9)
            obj_weights = aphal * self.obj_weights_init + (1 - aphal) * obj_weights
        return obj_weights
    
    def verb_weights_func(self, target_classes):
        if not self.verb_reweight:
            verb_weights = None
        elif self.use_static_weights:
            verb_weights = self.verb_weights_init
        else:
            verb_label_nums_in_batch = [0] * (self.num_verb_classes + 1)
            for target_class in target_classes:
                for label in target_class:
                    label_classes = torch.where(label > 0)[0]
                    if len(label_classes) == 0:
                        verb_label_nums_in_batch[-1] += 1
                    else:
                        for label_class in label_classes:
                            verb_label_nums_in_batch[label_class] += 1

            if self.q_verb.full(): self.q_verb.get()
            self.q_verb.put(np.array(verb_label_nums_in_batch))
            accumulated_verb_label_nums = np.sum(self.q_verb.queue, axis=0)
            verb_weights = self.cal_weights(accumulated_verb_label_nums, p=self.p_verb)

            aphal = min(math.pow(0.999, self.q_verb.qsize()),0.9)
            verb_weights = aphal * self.verb_weights_init + (1 - aphal) * verb_weights
        return verb_weights
